import os
from collections import defaultdict

from . import DATASET
from lib.base.base_dataset import BaseDataset


@DATASET.register('omniglot')
class Omniglot(BaseDataset):

    def __init__(self, cfg, type):
        super().__init__(cfg, type)

    def _get_categories(self):
        categories = []
        for path in ['images_background', 'images_evaluation']:
            for cate in os.listdir(os.path.join(self.data_path, path)):
                temp_path = os.path.join(self.data_path, path, cate)
                for each in os.listdir(temp_path):
                    categories.append(os.path.join(temp_path, each))
        return categories

    def _get_category2path(self) -> dict:
        category2path = defaultdict(list)
        categories = self._get_categories()
        for each in categories:
            dir_path = os.path.join(self.data_path, each)
            imgs_path = os.listdir(dir_path)
            imgs_path = [os.path.join(dir_path, x) for x in imgs_path]
            category2path[each].extend(imgs_path)
        return category2path
